# workstation_manager.py
import asyncio
import logging
from pathlib import Path
from typing import Dict, Optional, List
from uuid import uuid4
from pydantic import BaseModel
from browser_manager import BrowserManager, BrowserConfig
from callback_handler import CallbackHandler, CallbackConfig
from config_manager import config_manager
from logger_config import setup_logger
from service.db_service import SqlLiteService, DatabaseConfig
from exceptions import WorkstationError
from workstation.workstation import WorkstationConfig, Workstation


import logging

class ManagerConfig(BaseModel):
    """管理器全局配置"""
    max_workstations: int = 20
    health_check_interval: int = 5  # 秒
    auto_recovery: bool = True

    def __init__(self):
        super().__init__()
        try:
            self.max_workstations = config_manager.get_config("workstation_manager.max_workstations", self.max_workstations)
            self.health_check_interval = config_manager.get_config("workstation_manager.health_check_interval",
                                                           self.health_check_interval)
            self.auto_recovery = config_manager.get_config("workstation_manager.auto_recovery", self.auto_recovery)
        except KeyError as e:
            logging.error(f"Configuration key not found: {e}")
            # 使用默认值继续运行
        except ValueError as e:
            logging.error(f"Invalid configuration value: {e}")
            # 使用默认值继续运行
        except Exception as e:
            logging.error(f"Error loading config: {e}")
            raise e



class WorkstationManager:
    """工作台管理器（协调层）"""

    def __init__(
            self,
            manager_config: ManagerConfig = ManagerConfig(),
            browser_config: BrowserConfig = BrowserConfig(),
            db_config: DatabaseConfig = DatabaseConfig(),
            callback_config: CallbackConfig = CallbackConfig(),
            workstation_config: WorkstationConfig = WorkstationConfig(),
            logger: Optional[logging.Logger] = None
    ):
        # 初始化子系统
        self.browser_manager = BrowserManager(browser_config)
        self.db_service = SqlLiteService(db_config, logger=logger)
        self.callback_handler = CallbackHandler(callback_config)

        # 配置管理
        self.manager_config = manager_config
        self.workstation_config = workstation_config

        # 运行时状态
        self.workstations: Dict[str, Workstation] = {}
        self._health_check_task: Optional[asyncio.Task] = None
        self.logger = logger or logging.getLogger("workstation_manager")

    async def initialize(self):
        """异步初始化管理器"""
        await self.browser_manager.initialize()
        await self.db_service.initialize()
        await self.callback_handler.start()

        # 启动健康检查
        self._health_check_task = asyncio.create_task(self._monitor_health())
        self.logger.info("WorkstationManager initialized")

    async def create_workstation(
            self,
            name: str,
            workstation_id: str = None,
            initial_url: Optional[str] = None
    ) -> str:
        """创建新工作台"""
        if len(self.workstations) >= self.manager_config.max_workstations:
            raise WorkstationError("Maximum workstations reached")

        workstation_id = workstation_id or str(uuid4())

        # 这里需要先判断workstation_id是否已经在数据库存在
        workstation_data = await self.db_service.select_data(
            "workstations",
            condition="workstation_id = ?",
            params=(workstation_id,),
            single=True
        )
        # 如果存在则报错
        if workstation_data:
            raise WorkstationError(f"工作台ID({workstation_id})已经存在!")

        config = self.workstation_config.copy()
        if initial_url:
            config.initial_url = initial_url

        workstation = Workstation(
            workstation_id=workstation_id,
            browser_manager=self.browser_manager,
            db_service=self.db_service,
            callback_handler=self.callback_handler,
            config=config,
            logger=self.logger or setup_logger(f"workstation_{workstation_id}")
        )

        try:
            await workstation.ainit(name)
            self.workstations[workstation_id] = workstation
            # await self._persist_workstation(workstation_id, name)
            return workstation_id
        except Exception as e:
            await self._cleanup_failed_creation(workstation)
            raise WorkstationError(f"Creation failed: {str(e)}") from e

    async def open_workstation(self, workstation_id: str) -> str:
        """打开已存在的工作台"""
        try:
            # 检查是否已打开
            if workstation_id in self.workstations:
                raise WorkstationError(f"Workstation {workstation_id} is already open")

            # 从数据库获取配置
            record = await self.db_service.select_data(
                "workstations",
                condition="workstation_id = ?",
                params=(workstation_id,),
                single=True
            )

            if not record:
                raise WorkstationError(f"Workstation {workstation_id} not found in database")

            # 解析配置  兼容老版本
            if record["config"]:
                config = WorkstationConfig.parse_raw(record["config"])
            else:
                config = WorkstationConfig()
                config.cache_dir = Path(record["cache_dir"]) if record.get("cache_dir") else None

            # 创建工作站实例
            workstation = Workstation(
                workstation_id=workstation_id,
                browser_manager=self.browser_manager,
                db_service=self.db_service,
                callback_handler=self.callback_handler,
                config=config
            )


            # 使用持久化配置初始化
            await workstation.ainit()

            self.workstations[workstation_id] = workstation

            self.logger.info(f"Opened workstation {workstation_id}")
            return workstation_id

        except WorkstationError:
            raise

        except Exception as e:
             self.logger.error(f"Open workstation failed: {str(e)}")
             raise WorkstationError("Failed to open workstation") from e

    async def get_workstation(self, workstation_id: str) -> Workstation:
        """获取工作台实例"""
        if workstation := self.workstations.get(workstation_id):
            return workstation
        raise WorkstationError(f"Workstation {workstation_id} not found")

    async def close_workstation(
            self,
            workstation_id: str,
            remove_cache: bool = False
    ):
        """关闭指定工作台"""
        if workstation_id not in self.workstations:
            raise WorkstationError(f"Workstation {workstation_id} not found")

        workstation = self.workstations.pop(workstation_id)
        try:
            await workstation.aclose()
            self.logger.info(f"Closed workstation {workstation_id}")
        except Exception as e:
            self.logger.error(f"Error closing workstation: {str(e)}")
            raise WorkstationError("Close operation failed") from e

    async def close_all(self):
        """关闭所有工作台"""
        closing_tasks = [
            self.close_workstation(wid)
            for wid in list(self.workstations.keys())
        ]
        await asyncio.gather(*closing_tasks)
        self.logger.info("All workstations closed")

    async def list_workstations(
            self,
            keyword: str = None
    ) -> List[Dict]:
        """获取工作台列表"""
        self.logger.debug(f"List workstations: {keyword}")
        try:
            condition = "name LIKE ? OR workstation_id LIKE ?" if keyword else None
            params = ("%" + keyword + "%", "%" + keyword + "%") if keyword else None
            return await self.db_service.select_data(
                "workstations",
                condition=condition,
                params=params if params else None  # 适配数据库服务的参数要求
            )
        except Exception as e:
            self.logger.error(f"List workstations failed: {str(e)}")
            return []

    async def _persist_workstation(self, workstation_id: str, name: str = None, is_open: bool = True):
        """持久化工作台信息"""
        workstation_data = {
            "workstation_id": workstation_id,
            "is_open": is_open,
            "config": self.workstation_config.json()
        }
        if name:
            workstation_data["name"] = name
        try:
            await self.db_service.insert_data("workstations", workstation_data)
        except Exception as e:
            self.logger.error(f"Persist workstation failed: {str(e)}")
            await self.close_workstation(workstation_id)
            raise

    async def _cleanup_failed_creation(self, workstation: Workstation):
        """清理创建失败的资源"""
        try:
            await workstation.aclose()
            await self.browser_manager.close_context(
                workstation.workstation_id,
                remove_cache=True
            )
        except Exception as cleanup_error:
            self.logger.error(f"Cleanup failed: {str(cleanup_error)}")

    async def _monitor_health(self):
        """后台健康检查任务"""
        while True:
            try:
                await self._check_workstations_health()
                await asyncio.sleep(self.manager_config.health_check_interval)
            except Exception as e:
                self.logger.error(f"Health check failed: {str(e)}")

    async def _check_workstations_health(self):
        """执行健康检查"""
        for wid in list(self.workstations.keys()):
            workstation = self.workstations[wid]
            try:
                # 直接获取布尔值状态
                is_healthy = await workstation.health_check()
                if not is_healthy:
                    self.logger.warning(f"Restarting unhealthy workstation {wid}")
                    await self._recover_workstation(wid)
                    continue

                # 如果页面被关闭更新数据库
                if not is_healthy['page_active']:
                    self.logger.warning(f"Page inactive for workstation {wid}")
                    await self.close_workstation(wid)

            except Exception as e:
                self.logger.error(f"Health check error for {wid}: {str(e)}")
                await self.close_workstation(wid)

    async def _recover_workstation(self, workstation_id: str):
        """尝试恢复工作台"""
        if not self.manager_config.auto_recovery:
            return

        try:
            workstation = self.workstations[workstation_id]
            await workstation.aclose()
            await workstation.ainit()
            self.logger.info(f"Recovered workstation {workstation_id}")
        except Exception as e:
            self.logger.error(f"Recovery failed for {workstation_id}: {str(e)}")
            await self.close_workstation(workstation_id)

    async def shutdown(self):
        """完全关闭管理器"""
        try:
            # 有序关闭流程
            await self.close_all()
            await self.callback_handler.stop()
            await self.browser_manager.shutdown()
            await self.db_service.close()

            if self._health_check_task:
                self._health_check_task.cancel()

            self.logger.info("WorkstationManager shutdown complete")
        except Exception as e:
            self.logger.error(f"Shutdown error: {str(e)}")
            raise
